{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Mini Imagenet Generation\n",
    "\n",
    "This is one of a pair of notebooks used for generating an ImageNet-like dataset of training\n",
    "data using stable diffusion models. The difficulty of such artificial datasets can be\n",
    "easily tuned, and they are useful for debugging and testing deep learning applications.\n",
    "\n",
    "The first notebook uses Mistral-7B for taking class labels and generating descriptive prompts\n",
    "for image generation. The prompts are written out as shards to disk and shuffled. The process\n",
    "is parallelized using Ray.\n",
    "\n",
    "The second notebook uses Stable Diffustion to take descriptive prompts/image captions\n",
    "and renders them as image. This is a straightfowrard shard-to-shard transformation.\n",
    "\n",
    "Note that we are using explicit parallelization over shard files in the initial generation\n",
    "and the image generation, while we are using ray.data for the actual shuffling. That is\n",
    "because using explicit parallelization over shards makes it easier to restart jobs that have\n",
    "failed halfway through for some reason."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import matplotlib.pyplot as plt\n",
    "from pprint import pprint\n",
    "import webdataset as wds\n",
    "from diffusers import AutoPipelineForText2Image\n",
    "import torch\n",
    "import warnings\n",
    "import logging\n",
    "import logging\n",
    "import tqdm\n",
    "from IPython.display import display, clear_output\n",
    "from PIL import Image as PILImage\n",
    "from itertools import islice\n",
    "import glob\n",
    "import os\n",
    "import io\n",
    "from contextlib import contextmanager\n",
    "import sys\n",
    "\n",
    "\n",
    "class SuppressWarning:\n",
    "    def __enter__(self):\n",
    "        logging.disable(logging.WARNING)\n",
    "\n",
    "    def __exit__(self, type, value, traceback):\n",
    "        logging.disable(logging.NOTSET)\n",
    "\n",
    "\n",
    "tqdm.tqdm.disable = True\n",
    "\n",
    "def get_num_gpus():\n",
    "    cluster_resources = ray.cluster_resources()\n",
    "    return cluster_resources[\"GPU\"]\n",
    "\n",
    "@contextmanager\n",
    "def suppress_outputs(redirect):\n",
    "    old_stdout = sys.stdout\n",
    "    old_stderr = sys.stderr\n",
    "    \n",
    "    sys.stdout = redirect\n",
    "    sys.stderr = redirect\n",
    "    \n",
    "    try: \n",
    "        yield\n",
    "    finally:\n",
    "        sys.stdout = old_stdout\n",
    "        sys.stderr = old_stderr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# parameters\n",
    "\n",
    "odir = \"./mini-imagenet-10\"\n",
    "\n",
    "nactors = -1\n",
    "check_sufficient = True\n",
    "actor_startup_wait = 10"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Transformation Class\n",
    "\n",
    "We encapsulate the rendering into a `RenderPrompts` class. This class is instantiated once per GPU, loads the model, and then is ready to transform shards."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "class ShardTransformer:\n",
    "    def __init__(self):\n",
    "        self.pipe = AutoPipelineForText2Image.from_pretrained(\n",
    "            \"stabilityai/sdxl-turbo\", \n",
    "            torch_dtype=torch.float16, variant=\"fp16\"\n",
    "        ).to(\"cuda\")\n",
    "\n",
    "    def transform_shard(self, input_shard, output_shard, display_in_notebook=False):\n",
    "        ds = wds.WebDataset(input_shard).decode()\n",
    "        output = wds.TarWriter(output_shard)\n",
    "\n",
    "        for sample in ds:\n",
    "            sample = dict(sample)\n",
    "            text = sample[\"json\"][\"response\"]\n",
    "            with SuppressWarning():\n",
    "                image = self.pipe(prompt=text, num_inference_steps=4, guidance_scale=0.1).images[0]\n",
    "\n",
    "            sample[\"jpg\"] = image\n",
    "            output.write(sample)\n",
    "\n",
    "            if display_in_notebook:\n",
    "                clear_output(wait=True)\n",
    "                display(image)\n",
    "                pprint(text)\n",
    "\n",
    "        output.close()\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def maybe_clear_output():\n",
    "    try:\n",
    "        clear_output(wait=True)\n",
    "    except:\n",
    "        pass\n",
    "\n",
    "class RenderPrompts:\n",
    "    def __init__(self, display_in_notebook=False):\n",
    "        self.display_in_notebook = display_in_notebook\n",
    "\n",
    "    def gpu_is_sufficient(self):\n",
    "        return torch.cuda.get_device_properties(0).total_memory > 10**10\n",
    "\n",
    "    def load_model(self):\n",
    "        self.pipe = AutoPipelineForText2Image.from_pretrained(\n",
    "            \"stabilityai/sdxl-turbo\", \n",
    "            torch_dtype=torch.float16, variant=\"fp16\"\n",
    "        ).to(\"cuda\")\n",
    "\n",
    "    def transform_sample(self, sample):\n",
    "        text = sample[\"json\"][\"response\"]\n",
    "        with SuppressWarning():\n",
    "            image = self.pipe(prompt=text, num_inference_steps=4, guidance_scale=0.1).images[0]\n",
    "        sample[\"jpg\"] = image\n",
    "        return sample\n",
    "\n",
    "    def transform_sample_with_redirect(self, sample):\n",
    "        stdout = io.StringIO()\n",
    "        with suppress_outputs(stdout):\n",
    "            sample = self.transform_sample(sample)\n",
    "        sample[\"stdout\"] = stdout.getvalue()\n",
    "        return sample\n",
    "\n",
    "    def transform_shard(self, input_shard, output_shard, maxcount=999999999):\n",
    "        ds = wds.WebDataset(input_shard).decode()\n",
    "        output = wds.TarWriter(output_shard+\".temp\")\n",
    "\n",
    "        for sample in islice(ds, maxcount):\n",
    "            transformed_sample = self.transform_sample_with_redirect(dict(sample))\n",
    "            del transformed_sample[\"stdout\"]\n",
    "            maybe_clear_output()\n",
    "            output.write(transformed_sample)\n",
    "\n",
    "            if self.display_in_notebook:\n",
    "                clear_output(wait=True)\n",
    "                display(transformed_sample['jpg'])\n",
    "                pprint(transformed_sample[\"json\"][\"response\"])\n",
    "\n",
    "        output.close()\n",
    "\n",
    "        os.rename(output_shard+\".temp\", output_shard)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "transformer = RenderPrompts(display_in_notebook=True)\n",
    "transformer.load_model()\n",
    "shards = glob.glob(f\"{odir}/shuffled/*.tar\")\n",
    "transformer.transform_shard(shards[0], \"temp.tar\", maxcount=10)\n",
    "del transformer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Parallelization with Ray\n",
    "\n",
    "For parallel rendering, we use a Ray cluster. This will also work on a single machine with just one GPU."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import ray\n",
    "\n",
    "if not ray.is_initialized():\n",
    "    ray.init(log_to_driver=False)\n",
    "\n",
    "@ray.remote(num_gpus=1)\n",
    "class RayRenderPrompts(RenderPrompts):\n",
    "    def __init__(self):\n",
    "        super().__init__()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Start up and create the actor pool.\n",
    "# This tries to adapt to the number of GPUs available.\n",
    "# It also checks that each actor has sufficient memory.\n",
    "# If not, set up your cluster differently by excluding GPUs that are too small.\n",
    "# (Ray's facilities for heterogenous clusters are somewhat limited)\n",
    "\n",
    "ngpus = get_num_gpus() if nactors == -1 else nactors\n",
    "\n",
    "print(f\"using {ngpus} actors\")\n",
    "actors = [RayRenderPrompts.remote() for i in range(int(ngpus))]\n",
    "\n",
    "print(\"loading the models\")\n",
    "for actor in actors:\n",
    "    assert ray.get(actor.gpu_is_sufficient.remote()), \"GPU memory insufficient\"\n",
    "    ray.get(actor.load_model.remote())\n",
    "\n",
    "print(\"creating the pool\")\n",
    "pool = ray.util.ActorPool(actors)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import glob\n",
    "import os\n",
    "\n",
    "def apply_actor(actor, action):\n",
    "    src, dst = action\n",
    "    print(f\"START {src} -> {dst}\")\n",
    "    result = actor.transform_shard.remote(src, dst)\n",
    "    print(f\"DONE  {src} -> {dst}\")\n",
    "    return result\n",
    "\n",
    "!mkdir -p $odir/images\n",
    "shards = [os.path.basename(p) for p in  sorted(glob.glob(f\"{odir}/shuffled/*.tar\"))]\n",
    "actions = [(f\"{odir}/shuffled/{shard}\", f\"{odir}/images/{shard}\") for shard in shards]\n",
    "result = list(pool.map(apply_actor, actions))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
